import torch
from transformers import BertModel


# 定义了一个基于BERT的分类模型 继承自PyTorch的torch.nn.Module类
class BertClass(torch.nn.Module):
    def __init__(self):
        """
        初始化函数，用于定义模型结构和参数
        """
        super(BertClass, self).__init__()
        # 使用预训练的BERT模型加载权重，并作为模型的一部分。
        self.l1 = BertModel.from_pretrained("./models/")
        # 定义一个线性层，用于对BERT模型的输出进行线性变换
        self.pre_classifier = torch.nn.Linear(768, 768)
        # 定义一个Dropout层，用于在训练过程中进行随机失活
        self.dropout = torch.nn.Dropout(0.1)
        # 定义一个线性层，用于将BERT模型的输出映射到2个输出类别上
        self.classifier = torch.nn.Linear(768, 2)
        # 定义一个ReLU激活函数
        self.relu = torch.nn.ReLU()

    def forward(self, input_ids, attention_mask, token_type_ids):
        """
        定义前向传播函数，用于执行模型的前向计算
        """
        """
        在BERT的顶部添加两个线性层
        """
        # 通过预训练的BERT模型进行前向计算，得到BERT模型的输出
        output_1 = self.l1(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
        # 从BERT模型的输出中提取隐藏状态
        hidden_state = output_1[0]
        # 将隐藏状态的第一个词元的输出作为池化器（pooler）的输入
        pooler = hidden_state[:, 0]
        # 通过线性层进行线性变换
        pooler = self.pre_classifier(pooler)
        # 通过ReLU激活函数进行非线性变换
        pooler = self.relu(pooler)
        # 通过Dropout层进行随机失活
        pooler = self.dropout(pooler)
        # 通过线性层进行分类预测
        output = self.classifier(pooler)
        # 返回结果
        return output
